import requests
import pandas as pd
import os
from datetime import datetime

# 配置常量
API_URL = "http://10.223.1.1:16067/openapi/v1/space/worker_ip"
TOKEN = "f87b234e8d324c3ea78ba3390eafdeaf"
SPACE = "47"
WHITELIST_API = "https://10.223.12.1/openapi/v1.0/white-list"
API_KEY = "2c34ce91a4b6375d4e4a5c15944c15c3"
OPERATOR_ID = "d20f7105-ab6f-4e4c-99d8-ab1bdaad7258_1713422552"

SUCCESS_WEBHOOK = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=9e03ab9c-be12-4cf8-b395-00672690b2b7"
ERROR_WEBHOOK = "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=9e03ab9c-be12-4cf8-b395-00672690b2b7"
EXCEL_FILE = "/tools/py/last_ips.xlsx"
WHITELIST_HISTORY_FILE = "/tools/py/whitelist_history.xlsx"


def log(message, is_error=False):
    """格式化日志输出"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_type = "ERROR" if is_error else "INFO"
    print(f"[{timestamp}] [{log_type}] {message}")


def initialize_files():
    """初始化必要文件"""
    try:
        log("开始初始化文件...")
        for f in [EXCEL_FILE, WHITELIST_HISTORY_FILE]:
            if not os.path.exists(f):
                pd.DataFrame().to_excel(f, index=False)
                log(f"创建空文件: {f}")
        log("文件初始化完成")
    except Exception as e:
        log(f"文件初始化失败: {str(e)}", is_error=True)
        raise


def send_webhook(message, webhook_url=SUCCESS_WEBHOOK):
    """发送企业微信通知"""
    try:
        log(f"准备发送Webhook通知: {message[:50]}...")
        payload = {"msgtype": "text", "text": {"content": message}}
        response = requests.post(webhook_url, json=payload, timeout=5)
        response.raise_for_status()
        log("Webhook通知发送成功")
    except Exception as e:
        log(f"Webhook发送失败: {str(e)}", is_error=True)
        print(f"Webhook发送失败: {str(e)}")


def fetch_current_ips():
    """获取当前云图IP地址（带重试）"""
    try:
        log("正在获取云图IP地址...")
        response = requests.get(
            API_URL,
            headers={'TOKEN': TOKEN, 'space': SPACE},
            timeout=10
        )
        if response.status_code == 200:
            data = response.json()
            if data.get("code") == 200:
                ips = {item["internet_ip"] for item in data.get("data", [])}
                log(f"成功获取到{len(ips)}个IP地址")
                return ips
            else:
                log(f"API返回错误码: {data.get('code')}", is_error=True)
        else:
            log(f"HTTP状态码异常: {response.status_code}", is_error=True)
        return set()
    except Exception as e:
        log(f"获取IP失败: {str(e)}", is_error=True)
        return set()


def get_last_ips():
    """读取上次保存的IP集合"""
    try:
        log("正在读取历史IP文件...")
        if not os.path.exists(EXCEL_FILE):
            log("历史IP文件不存在，返回空集合")
            return set()

        df = pd.read_excel(EXCEL_FILE, engine='openpyxl')
        if 'ip' not in df.columns:
            log("历史文件缺少ip列，返回空集合")
            return set()

        ips = set(df['ip'])
        log(f"读取到{len(ips)}个历史IP")
        return ips
    except Exception as e:
        log(f"读取历史IP失败: {str(e)}", is_error=True)
        send_webhook(f"读取历史IP失败: {str(e)}", ERROR_WEBHOOK)
        return set()


def save_current_ips(ips):
    """保存当前IP到文件"""
    try:
        log(f"正在保存{len(ips)}个IP到文件")
        pd.DataFrame({'ip': list(ips)}).to_excel(EXCEL_FILE, index=False)
        log("IP保存成功")
    except Exception as e:
        log(f"保存IP失败: {str(e)}", is_error=True)
        send_webhook(f"保存IP失败: {str(e)}", ERROR_WEBHOOK)


def get_whitelist_history():
    """获取所有历史白名单ID"""
    try:
        log("正在读取白名单历史记录...")
        if not os.path.exists(WHITELIST_HISTORY_FILE):
            log("白名单历史文件不存在，返回空列表")
            return []

        df = pd.read_excel(WHITELIST_HISTORY_FILE, engine='openpyxl')
        if 'id' not in df.columns:
            log("历史文件缺少id列，返回空列表")
            return []

        ids = df['id'].tolist()
        log(f"读取到{len(ids)}条历史白名单记录")
        return ids
    except Exception as e:
        log(f"读取白名单历史失败: {str(e)}", is_error=True)
        send_webhook(f"读取白名单历史失败: {str(e)}", ERROR_WEBHOOK)
        return []


def update_whitelist_history(wh_id):
    """更新白名单历史记录"""
    try:
        log(f"正在更新白名单历史记录，新增ID: {wh_id}")
        history = get_whitelist_history()
        new_history = history + [wh_id]
        pd.DataFrame({'id': new_history}).to_excel(WHITELIST_HISTORY_FILE, index=False)
        log("白名单历史记录更新成功")
    except Exception as e:
        log(f"更新白名单历史失败: {str(e)}", is_error=True)
        send_webhook(f"更新白名单历史失败: {str(e)}", ERROR_WEBHOOK)


def delete_whitelist(wh_id):
    """删除指定白名单（带重试）"""
    try:
        log(f"正在删除白名单ID: {wh_id}")
        response = requests.delete(
            f"{WHITELIST_API}/{wh_id}",
            headers={"accept": "*/*", "apiKey": API_KEY},
            verify=False,
            timeout=10
        )
        if response.status_code == 200:
            res = response.json()
            if res.get("code") == 0:
                log(f"成功删除白名单ID: {wh_id}")
                return True
            else:
                log(f"删除失败，返回码: {res.get('code')}, 消息: {res.get('message')}", is_error=True)
        else:
            log(f"HTTP错误码: {response.status_code}", is_error=True)
        return False
    except Exception as e:
        log(f"删除白名单异常: {str(e)}", is_error=True)
        return False


def clear_old_whitelists():
    """清理所有历史白名单"""
    log("开始清理历史白名单")
    success_ids = []
    failed_ids = []
    history_ids = get_whitelist_history()

    for wh_id in history_ids:
        if delete_whitelist(wh_id):
            success_ids.append(wh_id)
        else:
            failed_ids.append(wh_id)

    log(f"清理完成: 成功{len(success_ids)}条, 失败{len(failed_ids)}条")

    # 更新历史记录
    try:
        if failed_ids:
            log(f"更新白名单历史记录，保留{len(failed_ids)}条失败记录")
            pd.DataFrame({'id': failed_ids}).to_excel(WHITELIST_HISTORY_FILE, index=False)
        else:
            log("无失败记录，清空白名单历史文件")
            pd.DataFrame().to_excel(WHITELIST_HISTORY_FILE, index=False)
    except Exception as e:
        log(f"更新历史记录失败: {str(e)}", is_error=True)

    return success_ids, failed_ids


def create_new_whitelist(ips):
    """创建新的白名单"""
    if not ips:
        log("无有效IP，跳过白名单创建")
        return None

    log(f"正在创建新白名单，包含{len(ips)}个IP")
    formatted_ips = '", "'.join(ips)
    payload = {
        # "endTimeStamp": 1747206000000,
        "filterCondition": f'srcAddress in ["{formatted_ips}"]',
        "ignore7DaysData": True,
        "name": "云图扫描动态IP",
        "open": True,
        "operatorId": OPERATOR_ID,
        "streamIdentifier": "ALL",
        "whiteListDesc": "自动化添加云图扫描动态IP"
    }

    try:
        response = requests.post(
            WHITELIST_API,
            json=payload,
            headers={
                "apiKey": API_KEY,
                "Content-Type": "application/json"
            },
            verify=False,
            timeout=15
        )
        if response.status_code == 200:
            res = response.json()
            if res.get("code") == 0:
                wh_id = res.get("data")
                log(f"成功创建白名单ID: {wh_id}")
                return wh_id
            else:
                log(f"创建失败，返回码: {res.get('code')}, 消息: {res.get('message')}", is_error=True)
        else:
            log(f"HTTP错误码: {response.status_code}", is_error=True)
        return None
    except Exception as e:
        log(f"创建白名单异常: {str(e)}", is_error=True)
        send_webhook(f"创建白名单异常: {str(e)}", ERROR_WEBHOOK)
        return None


def main_process():
    log("====== 开始执行主流程 ======")
    try:
        # 初始化文件
        initialize_files()

        # 获取IP信息
        log("正在获取当前IP地址...")
        current_ips = fetch_current_ips()
        if not current_ips:
            log("获取当前IP地址失败，终止流程", is_error=True)
            send_webhook("获取当前IP地址失败", ERROR_WEBHOOK)
            return

        log(f"当前获取到{len(current_ips)}个IP: {', '.join(list(current_ips)[:3])}...")

        # 读取历史IP

        last_ips = get_last_ips()
        log(f"历史IP数量: {len(last_ips)}")
        # print(last_ips)
        # print(type(last_ips))
        # print(current_ips)
        # print(type(current_ips))
        # 检测IP变化
        if current_ips != last_ips:
            log("检测到IP地址变化，开始更新白名单")

            # 清理旧白名单
            log("开始清理旧白名单...")
            success_del, failed_del = clear_old_whitelists()
            log_info = [
                f"清理结果: 成功{len(success_del)}条",
                f"失败{len(failed_del)}条" if failed_del else "全部清理成功"
            ]
            log(" | ".join(log_info))

            # 创建新白名单
            log("开始创建新白名单...")
            new_wh_id = create_new_whitelist(list(current_ips))

            if new_wh_id:
                # 更新记录
                log("开始更新白名单历史记录...")
                update_whitelist_history(new_wh_id)
                save_current_ips(current_ips)

                # 发送通知
                msg = [
                    "IP地址变更通知",
                    f"新增白名单ID: {new_wh_id}",
                    f"新增IP地址({len(current_ips)}个):",
                    "\n".join(current_ips)
                ]
                if success_del:
                    msg.insert(1, f"已清理历史白名单ID: {', '.join(map(str, success_del))}")
                if failed_del:
                    msg.insert(2, f"清理失败ID: {', '.join(map(str, failed_del))}")

                send_webhook("\n".join(msg))
                log("流程执行成功，已发送通知")
            else:
                log("白名单创建失败，流程终止", is_error=True)
                send_webhook("IP变更但白名单创建失败", ERROR_WEBHOOK)
        else:
            log("IP地址无变化，无需操作")
    except Exception as e:
        log(f"主流程执行异常: {str(e)}", is_error=True)
        send_webhook(f"流程执行异常: {str(e)}", ERROR_WEBHOOK)
    finally:
        log("====== 流程执行结束 ======\n")


if __name__ == "__main__":
    main_process()